fix(algorithms): raise instead of asserting on the skipped-reference-KL pairing - #3854
Open
tianyi-zhang-02 wants to merge 5 commits into
Open
fix(algorithms): raise instead of asserting on the skipped-reference-KL pairing#3854tianyi-zhang-02 wants to merge 5 commits into
tianyi-zhang-02 wants to merge 5 commits into
Conversation
…KL pairing
Sync PPO and sync GRPO both guard skip_reference_policy_logprobs_calculation
against a non-zero reference_policy_kl_penalty with a bare, message-less
assert master_config.loss_fn.reference_policy_kl_penalty == 0
python -O removes it entirely, so the run proceeds to train against a KL term
whose reference logprobs were never computed. Even with asserts enabled, a
message-less one says nothing about which two settings conflict.
async_ppo_train already does this as an if + ValueError. This is the same
defect class yuki-97 raised on NVIDIA-NeMo#3262 ("assert backend == 'vllm' gets stripped
under python -O"), which was converted there but left in these two siblings.
Both now use the async form, verbatim, so the three read alike.
Tests run under python -O as well as normally, which is what separates this
from a cosmetic change: with the assert restored, both modes fail.
Signed-off-by: Tianyi Zhang <zhangtianyi975@gmail.com>
Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>
…ed one grpo_train_sync's assert is unreachable in practice: grpo.setup checks the same pairing before the train loop starts, so on GRPO that is the one a user hits. It was an assert too, so under python -O both were stripped and nothing was left -- converting only the train-loop copy would have fixed the copy nobody reaches. PPO is the other way round: there is no setup-time equivalent, so the one in ppo_train is the only guard and converting it was already right. Keeps the grpo_sync conversion as well. It costs nothing and the three sites now read alike, which was the point. Test asserts on grpo.setup's source rather than calling it, since calling it needs a cluster. What it pins is that the reachable guard is not an assert. Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>
tianyi-zhang-02
force-pushed
the
fix-reference-kl-assert
branch
from
August 28, 2026 15:28
a3488d8 to
4bf8b1d
Compare
Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>
Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>
Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do?
Replaces skipped-reference-logprob configuration asserts in synchronous PPO and GRPO with
ValueErrors, so the safety check still runs underpython -O.The GRPO validation also runs at the start of setup, before logger, dataloader, checkpoint, or cluster initialization. The train-loop check remains as a defensive guard.
Validation
Final SHA:
d64af0e6d7dc4c29c24fe35df4ad2ae1a20d28c6, based on upstreammainatccbcd4cc5.tests/unit/algorithms/test_reference_kl_guard.py: 5 passedThis is setup/config validation, so no GPU workload is required.